"""
#===============================================================================
# License
#===============================================================================
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.

#===============================================================================
Copyright (c) 2009 bryan p. conrad
bconrad2@gmail.com

Repository Path:    
$HeadURL$

Created on:           Sep 28, 2009
Last Revised on:      $Date$
Version:              $Revision$
ID:                   $Id$
Author:               $Author$ :: (bryan)

#===============================================================================


#===============================================================================
# Documentation
#===============================================================================
This module can read in pyJT project files and animate the volumes using 
the specified kinematics.  It also provides a method for saving the 
each rendered pose as an image sequence.

This module can also be connected to an instance of pyJT to provide 3D
visualization of volume poses interactively.
"""

#===============================================================================
# Imports
#===============================================================================
import JT
import JT.Data2 as Data
import vtk
import time
import string
import sys
import os
import numpy as np
from PyQt4 import QtGui,QtCore
import threading

class Animate(object):
    '''
    classdocs
    '''


    def __init__(self):
        '''
        Constructor
        '''
        # This is probably the wrong place for this, i just want to have 
        # access to the standard QT dialogs (fileopen, color)
        self.app = QtGui.QApplication(sys.argv)
        
        self.drr = JT.AnimateDRR()
        self.drr.ren.SetBackground(0,0,0)
        style = vtk.vtkInteractorStyleTrackballCamera()
        style.AddObserver("CharEvent", self.CharEvent)
        self.drr.iren.SetInteractorStyle(style)
        self.data = Data.DataStructure()
        self.current_frame = 0
        
        # AxesActor
        axes = vtk.vtkAxesActor()
        axes.SetShaftTypeToCylinder()
        #axes.SetTotalLength(25, 25, 25)
        # Orientation marker to be used with vtkAxesActor.
        self.marker = vtk.vtkOrientationMarkerWidget()
        self.marker.SetOrientationMarker(axes)
        #self.marker.SetViewport(0,0,1,1)
        self.marker.SetViewport(0,0,.25,.25)
        
        self.marker.SetInteractor(self.drr.iren)
        self.marker.EnabledOn()
        self.marker.InteractiveOff()

        
#    def SetFileName(self,filename):
#        """
#        Read xml file to get data
#        """
#        self.data.load(filename)
    
    def CharEvent(self,object,event):
        key = self.drr.iren.GetKeySym()
        shift = self.drr.iren.GetShiftKey()
        ctrl = self.drr.iren.GetControlKey()
        #print key
        if key == 'p':
            self.Play()
        elif key == 'r':
            self.drr.ren.ResetCamera()
            self.drr.Render()
        elif key == 'o':
            self.Open()
        elif key == 's':
            self.Save()
        elif key == 'c':
            self.ChooseColor()
        elif key == 't':
            self.SetOpacityThreshold(ctrl)
        elif key == 'l':
            self.SetOpacityLevel(ctrl)
        elif (key in string.digits):
            self.ToggleVisibility(int(key)-1)
        elif key == 'comma':
            if ctrl:
                # Go to the beginning
                self.SetCurrentFrame(0)
            else:
                self.Backward()
        elif key == 'period':
            if ctrl:
                self.SetCurrentFrame(-1)
            else:
                self.Forward()
    
    def Backward(self):
        """
        Move forward or backward one frame
        """
        if (self.current_frame == 0):
            # Can't go back from 0!
            return
        self.current_frame -= 1
        i = self.current_frame
        print "Pose %i" % (i)
        for v,volume in enumerate(self.volume_list):
            self.volume_list[v].SetPose(self.poses[i,v,:])
        self.drr.Render()
    
    def Forward(self):
        """
        Move forward or backward one frame
        """
        if (self.current_frame==(len(self.poses)-1)):
            # Can't go forward from end!
            return
        self.current_frame += 1
        i = self.current_frame
        print "Pose %i" % (i)
        for v,volume in enumerate(self.volume_list):
            self.volume_list[v].SetPose(self.poses[i,v,:])
        self.drr.Render()
    
    def ToggleVisibility(self,index):
        try:
            volume = self.volume_list[index]
        except IndexError:
            print "No volume at that index"
            return
        if volume.GetVisibility():
            volume.VisibilityOff()
            self.drr.Render()
        else:
            volume.VisibilityOn()
            self.drr.Render()
    
    def Play(self):
        """
        Play back pose data
        """
        model_axes = vtk.vtkAxesActor()
        self.drr.ren.AddActor(model_axes)
        model_axes.SetTotalLength(1.5, 1.5, 1.5)
        
        for i in range(len(self.poses)):
            print "Playing Frame: %i" % i
            for v,volume in enumerate(self.volume_list):
                self.volume_list[v].SetPose(self.poses[i,v,:])
            model_axes.SetPosition(self.volume_list[1].GetPosition())
            model_axes.SetOrientation(self.volume_list[1].GetOrientation())
            time.sleep(.15)
            self.drr.Render()
        
    def LoadData(self):
        # Get list of images
        self.image_path_list =  self.data.get_path_list('FixedImage')
        # Get list of unique volume paths
        self.volume_path_list = []
        vol_list = self.data.get_path_list('Volume')
        for path in vol_list:
            if path not in self.volume_path_list:
                self.volume_path_list.append(path)
        
        print self.volume_path_list
        self.CreateVolumes(self.volume_path_list)
        
        self.cal = self.data.get_current_cal()
        self.poses = self.data.get_all_poses()
     
        self.drr.renWin.SetSize(300,300)
        self.drr.LoadCalFile(self.cal)
    
    def UpdatePoses(self,poses):
        self.poses = poses
    
    def SetCurrentFrame(self,index):
        if len(self.volume_list) == 0:
            # No volumes, cannot do anything
            return
        if index < 0:
            # Allow negative indices to start from the last frame
            index = (self.poses - 1 -index)
        self.current_frame = index
        i = index
        for v,volume in enumerate(self.volume_list):
            #print "Pose %i" % (i)
            self.volume_list[v].SetPose(self.poses[i,v,:])
            # Display Settings
        #self.drr.ren.ResetCamera()
        self.drr.Render()
        
    
    def CreateVolumes(self,volume_path_list):
        """
        """
        try:
            threshold,level = self.data.get_global_display_settings()
        except AttributeError:
            threshold,level = (60,0.15)
        self.volume_list = []
        self.drr.RemoveAllVolumes()
        for path in volume_path_list:
            ext = path[path.rfind('.')+1:].lower()
            if ext == 'stl':
                volume = JT.SurfaceModel()
                volume.LoadSTL(path)
            else:
                volume = JT.Volume2()
                volume.SetVolumeFileName(path)
                volume.SetOpacityThreshold(threshold)
                volume.SetOpacityLevel(level)
                center = np.asarray(volume.GetCenter()) -\
                         np.asarray(volume.GetPosition())
                volume.SetOrigin(center)
                # model might be a surface not a volume
                
            volume.SetOrientation(0,0,0)
            self.volume_list.append(volume)
            self.drr.AddVolume(volume)
            self.drr.AddAxesAtModelOrigin(volume)
            
    def Open(self,filename=None):
        """
        If no filename is given, a dialog will pop up to select one.
        """
        if filename==None: 
            self.filename = str(QtGui.QFileDialog.getOpenFileName())
        else:
            self.filename = str(filename)
        #self.SetFileName(self.filename)
        self.data.load(self.filename)
        self.LoadData()
        self.drr.ren.ResetCamera()
        self.drr.Render()
        self.drr.renWin.SetWindowName(os.path.basename(self.filename))
    
    def ChooseColor(self):
        color = QtGui.QColorDialog.getColor()
        rgb = color.getRgb()[0:3]
        for v,volume in enumerate(self.volume_list):
            self.volume_list[v].SetColor(rgb)
        self.drr.Render()
    
    def SetOpacityLevel(self,ctrl=False):
        """
        Set the color transfer function to alter display of image
        """
        point = self.volume_list[0]._otfLevel
        if (ctrl==False):
            point += 1.0/50
        else:
            point -= 1.0/50
        #don't do anything if point is less than 1 or greater than 254
        if 0< point <1:
            for v,volume in enumerate(self.volume_list):
                self.volume_list[v].SetOpacityLevel(point)
            self.drr.Render()
            print "Current Opacity Level: %f" % (point)

    def SetOpacityThreshold(self,ctrl=False):
        """
        Set the opacity transfer function to alter display of image
        """
        point = self.volume_list[0]._otfThreshold
        if (ctrl==False):
            point += 3.0
        else:
            point -= 3.0
        #don't do anything if point is less than 1 or greater than 254
        if 1< point < 254:
            for v,volume in enumerate(self.volume_list):
                self.volume_list[v].SetOpacityThreshold(point)
            self.drr.Render()
            print "Current Opacity Threshold: %f" % (point)
        
    def Save(self):
        """
        This will iterate through all the poses in the data structure and 
        save the
        """
        print "Saving Sequence..."
        basename = os.path.basename(self.filename).split('.')[0]
        dir = os.path.dirname(self.filename)
        try:
            os.mkdir(os.path.join(dir,'sequence'))
        except:
            pass
        for i in range(len(self.poses)):
            for v,volume in enumerate(self.volume_list):
                self.volume_list[v].SetPose(self.poses[i,v,:])
            filename = os.path.join(dir,'sequence',basename+"_"+str(i)+'.png')
            print "Saving: %s" % filename
            image = self.drr.SaveGeneratedImage(filename)
            
            self.drr.Render()

    def Test(self):
        filename = r'C:\Users\bryan\phd\data\Spine2\Spine2_C6C7_markers.xml'
        self.Open(filename)
        self.LoadData()
        for v,volume in enumerate(self.volume_list):
            self.volume_list[v].SetPose(self.poses[0,v,:])
        self.drr.ren.ResetCamera()

        self.drr.Render()
        self.drr.iren.Start()

    def Main(self):
        """
        Call this function to 
        """
        self.Open()
        self.drr.ren.ResetCamera()
        self.drr.Render()
        self.drr.iren.Start()

class ThreadedAnimate(Animate,threading.Thread):
    def __init__(self,filename):
        threading.Thread.__init__(self)
        Animate.__init__(self)
        self.filename = filename
    
    def run(self):
        self.Open(self.filename)
        self.LoadData()
        self.drr.Render()
        self.drr.iren.Start()
        print "Done"
        

class QTAnimate(Animate):
    '''
    classdocs
    '''
    def __init__(self):
        '''
        Constructor
        '''
        self.drrFrame = QtGui.QFrame()
        self.drr = JT.QTAnimateDRR(self.drrFrame)
        app.exec_()
        style = vtk.vtkInteractorStyleTrackballCamera()
        style.AddObserver("CharEvent", self.CharEvent)
        self.drr.iren.SetInteractorStyle(style)
        self.data = Data.DataStructure()
        self.current_frame = 0
        
    def Main(self):
        
        mainWindow = QtGui.QMainWindow()
        #frameWidget = QtGui.QFrame(mainWindow)
    
        mainWindow.setCentralWidget(self.drrFrame)
        #drr = JT.DRR.QTAnimateDRR(frameWidget)
        filename = r'C:\Users\bryan\phd\data\Spine1_8\test.xml'
        self.SetFileName(filename)
        self.LoadData()
        #print "Poses: %s" % str(self.poses[0,0,:])
        for v,volume in enumerate(self.volume_list):
            self.volume_list[v].SetPose(self.poses[0,v,:])

        self.drr.ren.ResetCamera()
        #self.drr.Render()
        #self.drr.iren.Start()
        
        mainWindow.show() 
         
class AnimateSTL(Animate):
    def CreateVolumes(self,model_path_list):
        """
        """
        self.volume_list = []
        self.drr.RemoveAllVolumes()
        for path in model_path_list:
            volume = JT.SurfaceModel()
            volume.LoadSTL(path)
            self.volume_list.append(volume)
            self.drr.AddVolume(volume)
    
    
    
def test_thread():
    filename = r'C:\Users\bryan\phd\data\Spine1_8\test.xml'
    t = ThreadedAnimate(filename)
    t.start()
    
    for i in range(46):
        t.Forward()
        time.sleep(.5)

def test_stl(filename=None):
    animate = Animate()
    if filename == None:
        #filename = r'C:\Users\bryan\Python\test_stl.xml'
        filename = r'C:\Users\bryan\phd\data\Spine2\Spine2_2\Spine2_2_C2-C7_test.xml'
    animate.Open(filename)
    animate.LoadData()
    animate.drr.ren.ResetCamera()

    animate.drr.Render()
    animate.drr.iren.Start()
    
    

if __name__ == '__main__':
    #animate = Animate()
#    animate.Test()
    test_stl()

        